{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "dd7c4c87",
   "metadata": {},
   "source": [
    "# Lesson 1: Router Engine"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "877df988",
   "metadata": {},
   "source": [
    "Welcome to Lesson 1.\n",
    "\n",
    "To access the `requirements.txt` file, the data/pdf file required for this lesson and the `helper` and `utils` modules, please go to the `File` menu and select`Open...`.\n",
    "\n",
    "I hope you enjoy this course!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "357c97c9",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab5f4f4a-5890-451c-8869-24606ef9f396",
   "metadata": {
    "height": 63,
    "tags": []
   },
   "outputs": [],
   "source": [
    "from helper import get_openai_api_key\n",
    "\n",
    "OPENAI_API_KEY = get_openai_api_key()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ffc9b4f4-64d4-4266-9889-54db90e00ee9",
   "metadata": {
    "height": 63,
    "tags": []
   },
   "outputs": [],
   "source": [
    "import nest_asyncio\n",
    "\n",
    "nest_asyncio.apply()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49fca250",
   "metadata": {},
   "source": [
    "## Load Data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f3ae2a8c",
   "metadata": {
    "tags": []
   },
   "source": [
    "To download this paper, below is the needed code:\n",
    "\n",
    "#!wget \"https://openreview.net/pdf?id=VtmBAGCN7o\" -O metagpt.pdf\n",
    "\n",
    "**Note**: The pdf file is included with this lesson. To access it, go to the `File` menu and select`Open...`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c7f012d-dcd3-4881-a568-72dd27d79159",
   "metadata": {
    "height": 97,
    "tags": []
   },
   "outputs": [],
   "source": [
    "from llama_index.core import SimpleDirectoryReader\n",
    "\n",
    "# load documents\n",
    "documents = SimpleDirectoryReader(input_files=[\"metagpt.pdf\"]).load_data()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b48a301",
   "metadata": {},
   "source": [
    "## Define LLM and Embedding model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a537bc0-78ee-4dda-a43f-60fd80062df6",
   "metadata": {
    "height": 80,
    "tags": []
   },
   "outputs": [],
   "source": [
    "from llama_index.core.node_parser import SentenceSplitter\n",
    "\n",
    "splitter = SentenceSplitter(chunk_size=1024)\n",
    "nodes = splitter.get_nodes_from_documents(documents)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de0660ee-b231-4351-b158-d8ad023e00b5",
   "metadata": {
    "height": 131,
    "tags": []
   },
   "outputs": [],
   "source": [
    "from llama_index.core import Settings\n",
    "from llama_index.llms.openai import OpenAI\n",
    "from llama_index.embeddings.openai import OpenAIEmbedding\n",
    "\n",
    "Settings.llm = OpenAI(model=\"gpt-3.5-turbo\")\n",
    "Settings.embed_model = OpenAIEmbedding(model=\"text-embedding-ada-002\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "997c7559",
   "metadata": {},
   "source": [
    "## Define Summary Index and Vector Index over the Same Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73d01b01-bc74-432a-8d92-07b9e86498b0",
   "metadata": {
    "height": 80,
    "tags": []
   },
   "outputs": [],
   "source": [
    "from llama_index.core import SummaryIndex, VectorStoreIndex\n",
    "\n",
    "summary_index = SummaryIndex(nodes)\n",
    "vector_index = VectorStoreIndex(nodes)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9898d3f",
   "metadata": {},
   "source": [
    "## Define Query Engines and Set Metadata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44cd7046-c714-4920-b077-b3ded917862f",
   "metadata": {
    "height": 97,
    "tags": []
   },
   "outputs": [],
   "source": [
    "summary_query_engine = summary_index.as_query_engine(\n",
    "    response_mode=\"tree_summarize\",\n",
    "    use_async=True,\n",
    ")\n",
    "vector_query_engine = vector_index.as_query_engine()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a1d6d75-247e-426a-8ef4-b49225c24796",
   "metadata": {
    "height": 301,
    "tags": []
   },
   "outputs": [],
   "source": [
    "from llama_index.core.tools import QueryEngineTool\n",
    "\n",
    "\n",
    "summary_tool = QueryEngineTool.from_defaults(\n",
    "    query_engine=summary_query_engine,\n",
    "    description=(\n",
    "        \"Useful for summarization questions related to MetaGPT\"\n",
    "    ),\n",
    ")\n",
    "\n",
    "vector_tool = QueryEngineTool.from_defaults(\n",
    "    query_engine=vector_query_engine,\n",
    "    description=(\n",
    "        \"Useful for retrieving specific context from the MetaGPT paper.\"\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98d2c152",
   "metadata": {},
   "source": [
    "## Define Router Query Engine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00734d7c-638a-4d63-ab1f-7f5a92a65119",
   "metadata": {
    "height": 233,
    "tags": []
   },
   "outputs": [],
   "source": [
    "from llama_index.core.query_engine.router_query_engine import RouterQueryEngine\n",
    "from llama_index.core.selectors import LLMSingleSelector\n",
    "\n",
    "\n",
    "query_engine = RouterQueryEngine(\n",
    "    selector=LLMSingleSelector.from_defaults(),\n",
    "    query_engine_tools=[\n",
    "        summary_tool,\n",
    "        vector_tool,\n",
    "    ],\n",
    "    verbose=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe3f0a76-68a8-444d-867f-d084bb3ff112",
   "metadata": {
    "height": 63,
    "tags": []
   },
   "outputs": [],
   "source": [
    "response = query_engine.query(\"What is the summary of the document?\")\n",
    "print(str(response))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3fedea0-f2a9-46bb-8aaf-287df65b8fff",
   "metadata": {
    "height": 29,
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(len(response.source_nodes))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af8c31b3-8e22-4ad9-9825-b8de21bd03c0",
   "metadata": {
    "height": 80,
    "tags": []
   },
   "outputs": [],
   "source": [
    "response = query_engine.query(\n",
    "    \"How do agents share information with other agents?\"\n",
    ")\n",
    "print(str(response))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aed060ee",
   "metadata": {},
   "source": [
    "## Let's put everything together"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8f92e0b-1c54-489b-b8dd-41ebaafb380a",
   "metadata": {
    "height": 63
   },
   "outputs": [],
   "source": [
    "from utils import get_router_query_engine\n",
    "\n",
    "query_engine = get_router_query_engine(\"metagpt.pdf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec1a43f3-77dc-472a-8adc-56551c00a0ff",
   "metadata": {
    "height": 63,
    "tags": []
   },
   "outputs": [],
   "source": [
    "response = query_engine.query(\"Tell me about the ablation study results?\")\n",
    "print(str(response))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
